"""
Created on 2025-01-10
@author: lihz
www.runoob.compytorchpytorch-tutorial.html
"""
import numpy as np
import torch
from matplotlib import pyplot as plt

"""
在PyTorch中，线性回归模型可以通过继承nn.Module类来实现。
"""

# =======================
# 【数据准备】生成一些符合简单线性关系的假数据，每个样本中有两个特征x1和x2
# =======================

torch.manual_seed(42)
X = torch.rand(100, 2)  # 生成1000个样本，每个样本有2个特征
true_w = torch.tensor([2.0, 3.0]).reshape(-1, 1)  # 真实权重向量，用于生成数据
true_b = 4.0  # 真实偏置项
# 解释：@ 符号表示矩阵乘法。(2,)的true_w广播为(2,1)的形式，然后进行矩阵乘法。
Y = X @ true_w + true_b + torch.randn(100).reshape(-1, 1) * 0.1
print(X.shape, Y.shape)

# =======================
# 【定义线性回归模型】
# 我们可以通过继承 nn.Module 来定义一个简单的线性回归模型。
# 在 PyTorch 中，线性回归的核心是 nn.Linear() 层，它会自动处理权重和偏置的初始化。
# =======================

import torch.nn as nn


class LinearRegressionModel(nn.Module):
    def __init__(self):
        super(LinearRegressionModel, self).__init__()  # 调用父类nn.Module的构造函数
        self.liner = nn.Linear(2, 1)  # 定义一个线性层，输入维度为2，输出维度为1

    def forward(self, x):
        out = self.liner(x)  # 前向传播，将输入x输入线性层，得到输出out
        return out


# =======================
# 【定义损失函数与优化器】
# 损失函数：线性回归的常见损失函数是 均方误差损失（MSELoss），用于衡量预测值与真实值之间的差异。PyTorch 中提供了现成的 MSELoss 函数。
# 梯度下降算法（优化器）：SGD（随机梯度下降） 或 Adam 优化器来最小化损失函数。
# =======================
model = LinearRegressionModel()
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.005)  # 梯度下降算法优化器需要与模型参数一起初始化

# =======================
# 【训练模型】
# 1. 使用输入数据 X 进行前向传播，得到预测值。
# 2. 计算损失（预测值与实际值之间的差异）。
# 3. 使用反向传播计算梯度。
# 4. 更新模型参数（权重和偏置）。
# =======================
num_epochs = 10000
for epoch in range(num_epochs):
    model.train()

    # 前向传播
    prediction = model(X)
    loss = criterion(prediction, Y)  # 计算损失(注意：预测值需要压缩为1维)

    # 反向传播和参数更新
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 打印损失
    if (epoch + 1) % 100 == 0:
        print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

# 查看训练后的权重和偏置
print('Model Weights:', model.liner.weight.data.numpy())
print('Model Bias:', model.liner.bias.data.numpy())

# =======================
# 【模型评估】因为没有加偏置项，所以图像显示没有完全重合，但基本是重合。
# =======================
with torch.no_grad():
    predictions = model(X)
    loss = criterion(predictions, Y)
    print(f'Test Loss: {loss.item():.4f}')

plt.scatter(X[:, 0], Y, color='blue', label='True Value')
plt.scatter(X[:, 0], predictions, color='red', label='Predicted Value')
plt.legend()
# plt.show()
plt.savefig('a-linear_regression.png')